package cn.trigram.concurrent.pool;

import cn.trigram.concurrent.ThreadPoolUtil;
import java.util.HashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import lombok.extern.slf4j.Slf4j;


/**
 * 基于权重值的线程池实现
 *
 * @version 1.0.0
 * @since 1.0.0
 */
public final class WeightThreadPoolExecutor extends ThreadPoolExecutor {

  /**
   * 这个数值表示已经提交的任务但是还没有完成的。这包括队列中的任务和已移交给Worker线程但后者尚未开始执行任务的任务。 这个值总是大于或等于 {@link #getActiveCount()}
   */
  private final AtomicInteger submittedCount = new AtomicInteger(0);

  public WeightThreadPoolExecutor(
      int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
      WeightQueue workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler
  ) {

    super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    workQueue.setExecutor(this);
    prestartAllCoreThreads(); // NOSOANR
  }

  @Override
  public void execute(Runnable command) {

    if (!(command instanceof WeightRunnable)) {
      throw new IllegalArgumentException("必须是WeightRunnable实例");
    }
    try {
      this.submittedCount.incrementAndGet();
      super.execute(command);
    } catch (RejectedExecutionException | NullPointerException e) {
      this.submittedCount.decrementAndGet();
      throw e;
    }
  }

  @Override
  protected void afterExecute(Runnable r, Throwable t) {
    //执行完成后应该把submittedCount减一
    submittedCount.decrementAndGet();
  }

  public int getSubmittedCount() {

    return submittedCount.get();
  }

  @Slf4j
  protected static class WeightQueue extends LinkedBlockingQueue<Runnable> {

    /**
     * 容量
     */
    private final int capacity;

    private transient volatile WeightThreadPoolExecutor executor = null;

    /**
     * 保留容量阈值比例；0f-1f
     */
    private final float reservedCapacityThresholdRatio;

    /**
     * @param capacity
     *     容量
     */
    public WeightQueue(int capacity) {

      this(capacity, 0f);
    }

    /**
     * @param capacity
     *     容量
     * @param reservedCapacityThresholdRatio
     *     保留容量阈值比例
     */
    public WeightQueue(int capacity, float reservedCapacityThresholdRatio) {

      super(capacity);
      if (!(reservedCapacityThresholdRatio >= 0f && reservedCapacityThresholdRatio <= 1f)) {
        throw new IllegalArgumentException("reservedCapacityThresholdRatio保留阈值比例应该在[0f-1f]之间");
      }
      this.capacity                       = capacity;
      this.reservedCapacityThresholdRatio = reservedCapacityThresholdRatio;
    }

    public void setExecutor(WeightThreadPoolExecutor tp) {

      executor = tp;
    }

    public boolean force(Runnable o) {

      if (executor.isShutdown()) {
        throw new RejectedExecutionException("Executor not running, can't force a command into the queue");
      }
      return super.offer(o);
    }

    public boolean force(Runnable o, long timeout, TimeUnit unit) throws InterruptedException {

      if (executor.isShutdown()) {
        throw new RejectedExecutionException("Executor not running, can't force a command into the queue");
      }
      return super.offer(o, timeout, unit);
    }

    /**
     * 提交任务，有以下判断顺序：
     * <ol>
     *   <li>如果剩余容量小于了设置的剩余容量限值，此时返回的false表示应该执行拒绝策略</li>
     *   <li>当前可用线程数大于等于最大线程数，任务提交到队列中</li>
     *   <li>线程池中的已提交数小于当前可用线程数，代表有空闲的线程，但是也是提交到队列等待线程自动从队列获取任务处理</li>
     *   <li>当前可用线程数小于最大线程数时，此时返回的false是表示应该直接创建一个新的线程</li>
     *   <li>其它情况直接提交队列</li>
     * </ol>
     */
    @Override
    public boolean offer(Runnable runnable) {
      //如果剩余容量小于了设置的剩余容量限值，此时返回的false表示应该执行拒绝策略
      if (this.remainingCapacity() < Math.ceil(this.capacity * reservedCapacityThresholdRatio)) {
        return false;
      }
      // springside: threadPool.getPoolSize() 是个有锁的操作，所以尽量减少
      int currentPoolSize = executor.getPoolSize();
      // 当前可用线程数大于等于最大线程数，任务提交到队列中
      if (currentPoolSize >= executor.getMaximumPoolSize()) {
        return super.offer(runnable);
      }
      // 有空闲线程并且当前线程数已经达到最大线程数，任务提交到队列中
      if (executor.getSubmittedCount() < currentPoolSize) {
        return super.offer(runnable);
      }
      // 当前线程数少于最大线程数，强制创建线程
      if (currentPoolSize < executor.getMaximumPoolSize()) {
        return false;
      }
      // 最后的保险，直接提交队列
      return super.offer(runnable);
    }

    public void showStatsInfo() {

      final WeightRunnable[] runnables        = super.toArray(new WeightRunnable[super.size()]);
      final int              reservedCapacity = (int) Math.ceil(this.capacity * reservedCapacityThresholdRatio);

      log.debug("非保留容量区统计信息[权重：总数]：");
      HashMap<String, Integer> map = new HashMap<>();
      for (int i = 0; i < capacity - reservedCapacity; i++) {
        final WeightRunnable runnable = runnables[i];
        map.merge(String.valueOf(runnable.getWeight()), 1, Integer::sum);
      }
      log.debug("{}", map);
      log.debug("保留容量区统计信息[权重：总数]：");
      map = new HashMap<>();
      for (int i = capacity - reservedCapacity; i < capacity; i++) {
        final WeightRunnable runnable = runnables[i];
        map.merge(String.valueOf(runnable.getWeight()), 1, Integer::sum);
      }
      log.debug("{}", map);
    }

  }

  /**
   * 基于权重值的线程任务。 默认权重为{@link Integer#MAX_VALUE}。 数值越小，权重越大
   *
   * @version 1.0.0
   * @since 1.0.0
   */
  public static class WeightRunnable implements Runnable {

    /**
     * 权重值。数值越小，权重优先级越高；默认权重为{@link Integer#MAX_VALUE}。
     */
    private final int weight;

    private final Runnable runnable;

    public WeightRunnable(Runnable runnable) {

      this(Integer.MAX_VALUE, runnable);
    }

    public WeightRunnable(int weight, Runnable runnable) {

      this.weight   = weight;
      this.runnable = runnable;
    }

    @Override
    public void run() {

      runnable.run();
    }

    public int getWeight() {

      return weight;
    }

  }

  /**
   * 基于权重的拒绝策略
   *
   * @version 1.0.0
   * @since 1.0.0
   */
  @Slf4j
  public static class WeightPolicyRejectedExecutionHandler implements RejectedExecutionHandler {

    /**
     * 权重值的阈值。表示优先级高于这个阈值的权重才能直接强制进入队列的保留容量中。默认值{@link Integer#MIN_VALUE}
     */
    private final int weightThreshold;

    public WeightPolicyRejectedExecutionHandler() {

      this(Integer.MIN_VALUE);
    }

    public WeightPolicyRejectedExecutionHandler(int weightThreshold) {

      this.weightThreshold = weightThreshold;
    }

    @Override
    public void rejectedExecution(Runnable runnable, ThreadPoolExecutor executor) {

      if (log.isDebugEnabled()) {
        // 调试模式下打印线程池信息
        ((WeightQueue) executor.getQueue()).showStatsInfo();
        ThreadPoolUtil.recordStats(executor);
      }
      WeightRunnable weightRunnable = (WeightRunnable) runnable;
      final int      weight         = weightRunnable.getWeight();
      final boolean  judgeResult    = weight < weightThreshold;
      if (log.isDebugEnabled()) {
        log.debug("非保留区已满，当前任务的权重为[{}]，权重判定结果：{}", weight, judgeResult);
      }
      // 如果运行任务的权重值高于了给定的权重阈值，则直接强制提交
      if (judgeResult) {
        final boolean forceResult = ((WeightQueue) executor.getQueue()).force(weightRunnable);
        if (!forceResult && log.isDebugEnabled()) {
          log.debug("权重线程池队列已满");
        }
      }
    }

  }

}

